package com.csust.zerocode.ai;

import com.csust.zerocode.ai.chatmemory.RedisChatMemoryStoreWithDb;
import com.csust.zerocode.ai.guardrail.PromptSafetyInputGuardrail;
import com.csust.zerocode.ai.tools.*;
import com.csust.zerocode.exception.BusinessException;
import com.csust.zerocode.exception.ErrorCode;
import com.csust.zerocode.model.enums.CodeGenTypeEnum;
import com.csust.zerocode.service.ChatHistoryService;
import com.csust.zerocode.utils.SpringContextUtil;
import com.github.benmanes.caffeine.cache.Cache;
import com.github.benmanes.caffeine.cache.Caffeine;
import dev.langchain4j.data.message.ToolExecutionResultMessage;
import dev.langchain4j.memory.chat.MessageWindowChatMemory;
import dev.langchain4j.model.chat.ChatModel;
import dev.langchain4j.model.chat.StreamingChatModel;
import dev.langchain4j.service.AiServices;
import jakarta.annotation.Resource;
import lombok.extern.slf4j.Slf4j;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;

import java.time.Duration;

import static com.csust.zerocode.model.enums.CodeGenTypeEnum.*;

@Configuration
@Slf4j
public class AiCodeGeneratorServiceFactory {

    @Resource(name = "openAiChatModel")
    private ChatModel chatModel;
    @Resource
    private RedisChatMemoryStoreWithDb redisChatMemoryStoreDb;
    @Resource
    private ChatHistoryService chatHistoryService;
    @Resource
    private ToolManager toolManager;
    /**
     * AI 服务实例缓存
     * 缓存策略：
     * - 最大缓存 1000 个实例
     * - 写入后 30 分钟过期
     * - 访问后 10 分钟过期
     */
    private final Cache<String, AiCodeGeneratorService> serviceCache = Caffeine.newBuilder()
            .maximumSize(1000)
            .expireAfterWrite(Duration.ofMinutes(30))
            .expireAfterAccess(Duration.ofMinutes(10))
            .removalListener((key, value, cause) -> {
                log.debug("AI 服务实例被移除，cacheKey: {}, 原因: {}", key, cause);
            })
            .build();

    /**
     * 通过缓存获取AI服务实例
     * @param appId 应用ID
     * @return AI服务实例
     */
    public AiCodeGeneratorService getAiCodeGeneratorService(Long appId) {
        return getAiCodeGeneratorService(appId,HTML);
    }
    /**
     * 通过缓存获取AI服务实例
     * @param appId 应用ID
     * @return AI服务实例
     */
    public AiCodeGeneratorService getAiCodeGeneratorService(Long appId, CodeGenTypeEnum codeGenTypeEnum) {
        String cacheKey = buildCacheKey(appId, codeGenTypeEnum);
        return serviceCache.get(cacheKey,key->createAiCodeGeneratorService(appId,codeGenTypeEnum));
    }

    /**
     * 创建AI服务实例
     * @param appId 应用ID
     * @return AI服务实例
     */
    private AiCodeGeneratorService createAiCodeGeneratorService(Long appId, CodeGenTypeEnum codeGenTypeEnum) {
        //构建对话历史持久化容器
        MessageWindowChatMemory chatMemory = MessageWindowChatMemory
                .builder()
                .id(appId)
                .chatMemoryStore(redisChatMemoryStoreDb)
                .maxMessages(20)
                .build();
        // 从数据库加载历史对话到记忆中
        chatHistoryService.loadChatHistoryToMemory(appId, chatMemory, 20);
        return  switch (codeGenTypeEnum){
            case VUE_PROJECT -> {
                // 使用多例模式的 StreamingChatModel 解决并发问题
                StreamingChatModel reasoningStreamingChatModel = SpringContextUtil.getBean("reasoningStreamingChatModelPrototype", StreamingChatModel.class);
                yield AiServices.builder(AiCodeGeneratorService.class)
                        .streamingChatModel(reasoningStreamingChatModel)
                        .chatMemoryProvider(memoryId -> chatMemory)
                        .tools(toolManager.getAllTools())
                        .inputGuardrails(new PromptSafetyInputGuardrail())  // 添加输入护轨
                        //解决提示工具不存在 ai幻觉问题
                        .hallucinatedToolNameStrategy(toolExecutionRequest -> ToolExecutionResultMessage.from(
                                toolExecutionRequest, "Error: there is no tool called " + toolExecutionRequest.name()
                        ))
                        .build();
            }
            case HTML, MULTI_FILE -> {
                //设置 对话持久化，并且返回ai服务
                // 使用多例模式的 StreamingChatModel 解决并发问题
                StreamingChatModel openAiStreamingChatModel = SpringContextUtil.getBean("streamingChatModelPrototype", StreamingChatModel.class);
                yield AiServices.builder(AiCodeGeneratorService.class)
                        .chatModel(chatModel)
                        .chatMemory(chatMemory)
                        .streamingChatModel(openAiStreamingChatModel)
                        .inputGuardrails(new PromptSafetyInputGuardrail())  // 添加输入护轨
                        .build();
            }

        default -> throw new BusinessException(ErrorCode.SYSTEM_ERROR,
                "不支持的代码生成类型: " + codeGenTypeEnum.getValue());

        };

    }
    @Bean
    public AiCodeGeneratorService aiCodeGeneratorService() {
        return getAiCodeGeneratorService(0L);
    }


    private String buildCacheKey(Long appId, CodeGenTypeEnum codeGenTypeEnum){
        return appId + "_" + codeGenTypeEnum.getValue();
    }
}
